package udpserver

import (
	"fmt"
	"net"
	"sync"
	"time"
)

type UDPServer struct {
	addr        *net.UDPAddr
	conn        *net.UDPConn
	clientMap   map[string]*ClientStatus
	mu          sync.Mutex
	stopChan    chan bool
	receiveChan chan string
	timeout     time.Duration
}

type ClientStatus struct {
	LastHeartbeat   time.Time
	OnlineDuration  time.Duration
	TotalOnlineTime time.Duration
}

func NewUDPServer(port uint16, timeout time.Duration) (*UDPServer, error) {
	addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", port))
	if err != nil {
		return nil, err
	}

	server := &UDPServer{
		addr:        addr,
		clientMap:   make(map[string]*ClientStatus),
		stopChan:    make(chan bool),
		receiveChan: make(chan string),
		timeout:     timeout,
	}

	return server, nil
}

func (s *UDPServer) Start() error {
	conn, err := net.ListenUDP("udp", s.addr)
	if err != nil {
		return err
	}
	s.conn = conn
	go s.listen()
	go s.monitorClients()
	return nil
}

func (s *UDPServer) Stop() {
	s.stopChan <- true
}

func (s *UDPServer) Close() {
	if s.conn != nil {
		_ = s.conn.Close()
	}
	close(s.receiveChan)
	close(s.stopChan)
}

func (s *UDPServer) Receive() <-chan string {
	return s.receiveChan
}

func (s *UDPServer) listen() {
	buf := make([]byte, 1024)
	for {
		select {
		case <-s.stopChan:
			return
		default:
			n, _, err := s.conn.ReadFromUDP(buf)
			if err != nil {
				fmt.Println("Error reading from UDP:", err)
				continue
			}

			clientID := string(buf[:n])
			s.mu.Lock()
			status, exists := s.clientMap[clientID]
			if !exists {
				status = &ClientStatus{LastHeartbeat: time.Now()}
				s.clientMap[clientID] = status
			}
			status.LastHeartbeat = time.Now()
			s.mu.Unlock()

			s.receiveChan <- clientID
			//fmt.Println("Received go-base from:", clientID, "from", remoteAddr)
		}
	}
}

func (s *UDPServer) monitorClients() {
	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-s.stopChan:
			return
		case <-ticker.C:
			s.mu.Lock()
			currentTime := time.Now()

			for _, status := range s.clientMap {
				if currentTime.Sub(status.LastHeartbeat) > s.timeout {
					status.TotalOnlineTime += status.OnlineDuration
					status.OnlineDuration = 0
					//fmt.Printf("Client %s timed out. Total online time: %v\n", clientID, status.TotalOnlineTime)
				} else {
					status.OnlineDuration += 1 * time.Second
					//fmt.Printf("Client %s is online. Online duration: %v\n", clientID, status.OnlineDuration)
				}
			}
			s.mu.Unlock()
		}
	}
}
